from keras import Model
from keras.layers import Activation,Conv2D,MaxPool2D,UpSampling2D,Conv3D,MaxPool3D,concatenate,UpSampling3D,BatchNormalization,Input,Dropout
from keras.utils import plot_model

class XDUnet:    
    def unetConv2D(self,channels):
        return Conv2D(filters=channels,kernel_size=3,strides=1,padding='same',activation='relu',kernel_initializer='glorot_normal',bias_initializer='he_normal')

    def unetConv3D(self,channels):
        return Conv3D(filters=channels,kernel_size=[1,3,3],strides=1,padding='same',activation='relu',kernel_initializer='glorot_normal',bias_initializer='he_normal')

    def Unet2D(self,topChannel,n_classes,input_width,input_height):
        assert input_width%16 == 0,"输入宽度不合要求"
        assert input_height%16 == 0,"输入高度不合要求"

        img_input = Input(shape=(input_width,input_height,1))
        # 暂时先缓冲一下，直接1->topchannels梯度太大可能有不好结果
        # 1->10->topchannels
        temp    = self.unetConv2D(10)(img_input)

        conv1_1 = self.unetConv2D(topChannel)(temp)
        conv1_2 = self.unetConv2D(topChannel)(conv1_1)
        bn1     = BatchNormalization()(conv1_2)
        pool1   = MaxPool2D(strides=2,padding='same')(bn1)
        
        conv2_1 = self.unetConv2D(topChannel*2)(pool1)
        conv2_2 = self.unetConv2D(topChannel*2)(conv2_1)
        bn2     = BatchNormalization()(conv2_2)
        pool2   = MaxPool2D(strides=2,padding='same')(bn2)

        conv3_1 = self.unetConv2D(topChannel*4)(pool2)
        conv3_2 = self.unetConv2D(topChannel*4)(conv3_1)
        bn3     = BatchNormalization()(conv3_2)
        pool3   = MaxPool2D(strides=2,padding='same')(bn3)
        
        conv4_1 = self.unetConv2D(topChannel*8)(pool3)
        conv4_2 = self.unetConv2D(topChannel*8)(conv4_1)
        bn4     = BatchNormalization()(conv4_2)
        pool4   = MaxPool2D(strides=2,padding='same')(bn4)

        conv5_1 = self.unetConv2D(topChannel*16)(pool4)
        conv5_2 = self.unetConv2D(topChannel*16)(conv5_1)
        conv5_3 = self.unetConv2D(topChannel*16)(conv5_2)
        bn5     = BatchNormalization()(conv5_3)

        upsamp1 = UpSampling2D(size=(2,2))(bn5)
        conta1  = concatenate([upsamp1,bn4],axis=-1)
        ucon1_1 = self.unetConv2D(topChannel*8)(conta1)
        ucon1_2 = self.unetConv2D(topChannel*8)(ucon1_1)
        ubn1    = BatchNormalization()(ucon1_2)

        upsamp2 = UpSampling2D(size=(2,2))(ubn1)
        conta2  = concatenate([upsamp2,bn3],axis=-1)
        ucon2_1 = self.unetConv2D(topChannel*4)(conta2)
        ucon2_2 = self.unetConv2D(topChannel*4)(ucon2_1)
        ubn2    = BatchNormalization()(ucon2_2)

        upsamp3 = UpSampling2D(size=(2,2))(ubn2)
        conta3  = concatenate([upsamp3,bn2],axis=-1)
        ucon3_1 = self.unetConv2D(topChannel*2)(conta3)
        ucon3_2 = self.unetConv2D(topChannel*2)(ucon3_1)
        ubn3    = BatchNormalization()(ucon3_2)

        upsamp4 = UpSampling2D(size=(2,2))(ubn3)
        conta4  = concatenate([upsamp4,bn1],axis=-1)
        ucon4_1 = self.unetConv2D(topChannel)(conta4)
        ucon4_2 = self.unetConv2D(topChannel)(ucon4_1)

        out1    = Conv2D(filters=n_classes,kernel_size=1,strides=1,padding='same',activation='relu',kernel_initializer='glorot_normal',bias_initializer='he_normal')(ucon4_2)
        out2    = Dropout(0.5)(out1)
        out3    = Activation('softmax')(out2)

        model   = Model(inputs=img_input,outputs=out3)
        return model

    def Unet3D(self,topChannel,n_classes,input_width,input_height):
        assert input_width%16 == 0,"输入宽度不合要求"
        assert input_height%16 == 0,"输入高度不合要求"
                
        img_input = Input(shape=(None,input_width,input_height,1))

        temp    = self.unetConv3D(10)(img_input)
        
        conv1_1 = self.unetConv3D(topChannel)(temp)
        conv1_2 = self.unetConv3D(topChannel)(conv1_1)
        bn1     = BatchNormalization()(conv1_2)
        pool1   = MaxPool3D(strides=[1,2,2],pool_size=(1,2,2),padding='same')(bn1)
        
        conv2_1 = self.unetConv3D(topChannel*2)(pool1)
        conv2_2 = self.unetConv3D(topChannel*2)(conv2_1)
        bn2     = BatchNormalization()(conv2_2)
        pool2   = MaxPool3D(strides=[1,2,2],pool_size=(1,2,2),padding='same')(bn2)

        conv3_1 = self.unetConv3D(topChannel*4)(pool2)
        conv3_2 = self.unetConv3D(topChannel*4)(conv3_1)
        bn3     = BatchNormalization()(conv3_2)
        pool3   = MaxPool3D(strides=[1,2,2],pool_size=(1,2,2),padding='same')(bn3)
        
        conv4_1 = self.unetConv3D(topChannel*8)(pool3)
        conv4_2 = self.unetConv3D(topChannel*8)(conv4_1)
        bn4     = BatchNormalization()(conv4_2)
        pool4   = MaxPool3D(strides=[1,2,2],pool_size=(1,2,2),padding='same')(bn4)

        conv5_1 = self.unetConv3D(topChannel*16)(pool4)
        conv5_2 = self.unetConv3D(topChannel*16)(conv5_1)
        conv5_3 = self.unetConv3D(topChannel*16)(conv5_2)
        bn5     = BatchNormalization()(conv5_3)

        upsamp1 = UpSampling3D(size=(1,2,2))(bn5)
        conta1  = concatenate([upsamp1,bn4],axis=-1)
        ucon1_1 = self.unetConv3D(topChannel*8)(conta1)
        ucon1_2 = self.unetConv3D(topChannel*8)(ucon1_1)
        ubn1    = BatchNormalization()(ucon1_2)

        upsamp2 = UpSampling3D(size=(1,2,2))(ubn1)
        conta2  = concatenate([upsamp2,bn3],axis=-1)
        ucon2_1 = self.unetConv3D(topChannel*4)(conta2)
        ucon2_2 = self.unetConv3D(topChannel*4)(ucon2_1)
        ubn2    = BatchNormalization()(ucon2_2)

        upsamp3 = UpSampling3D(size=(1,2,2))(ubn2)
        conta3  = concatenate([upsamp3,bn2],axis=-1)
        ucon3_1 = self.unetConv3D(topChannel*2)(conta3)
        ucon3_2 = self.unetConv3D(topChannel*2)(ucon3_1)
        ubn3    = BatchNormalization()(ucon3_2)

        upsamp4 = UpSampling3D(size=(1,2,2))(ubn3)
        conta4  = concatenate([upsamp4,bn1],axis=-1)
        ucon4_1 = self.unetConv3D(topChannel)(conta4)
        ucon4_2 = self.unetConv3D(topChannel)(ucon4_1)

        out1    = Conv3D(filters=n_classes,kernel_size=1,strides=1,padding='same',activation='relu',kernel_initializer='glorot_normal',bias_initializer='he_normal')(ucon4_2)
        out2    = Dropout(out1)
        out3    = Activation('softmax')(out2)

        model   = Model(inputs=img_input,outputs=out3)
        return model

if __name__ == "__main__":
    Unet = XDUnet()
    model2D = Unet.Unet2D(64,2,320,320)
    model3D = Unet.Unet3D(64,7,64,64)
    plot_model(model2D, show_shapes=True, to_file='model_unet.png')